#include "perspective_camera_model.h"
#include "opencv2/core/eigen.hpp"
#include "opencv2/imgproc.hpp"
#include "opencv2/calib3d.hpp"

PerspectiveCameraModel::PerspectiveCameraModel(const Eigen::Matrix3d& K,
                                               const Sophus::SE3d& pose_world_camera,
                                               const Vector5d& distortion_coeffs)
  : K_{K}
  , pose_world_camera_{pose_world_camera}
  , distortion_coeffs_{distortion_coeffs}
{
  camera_projection_matrix_ = computeCameraProjectionMatrix();
}

Sophus::SE3d PerspectiveCameraModel::getPose() const
{
  return pose_world_camera_;
}

Eigen::Matrix3d PerspectiveCameraModel::getCalibrationMatrix() const
{
  return K_;
}

PerspectiveCameraModel::Matrix34d PerspectiveCameraModel::getCameraProjectionMatrix() const
{
  return camera_projection_matrix_;
}

Eigen::Vector2d PerspectiveCameraModel::projectWorldPoint(Eigen::Vector3d world_point) const
{
  // Todo: Implement projection using camera_projection_matrix_.
  // p_pixel = P * [X; Y; Z; 1]
  Eigen::Vector4d world_point_homogeneous(world_point.x(), world_point.y(), world_point.z(), 1.0);
  Eigen::Vector3d projected = camera_projection_matrix_ * world_point_homogeneous;

  // Normalize
  if (projected.z() <= 0.0){
    return Eigen::Vector2d::Constant(std::numeric_limits<double>::quiet_NaN());
  }
  return Eigen::Vector2d(projected.x() / projected.z(), projected.y() / projected.z());
}

Eigen::Matrix2Xd PerspectiveCameraModel::projectWorldPoints(Eigen::Matrix3Xd world_points) const
{
  // Todo: Optionally implement projection using camera_projection_matrix_.
  const int num_points = world_points.cols();
  Eigen::Matrix2Xd projected_points(2, num_points);
  for (int i = 0; i < num_points; ++i){
    projected_points.col(i) = projectWorldPoint(world_points.col(i));
  }
  return projected_points;
}

PerspectiveCameraModel::Matrix34d PerspectiveCameraModel::computeCameraProjectionMatrix()
{
  // Todo: Compute camera projection matrix.
  // Extract rotation and translation from the pose
  Eigen::Matrix3d R = pose_world_camera_.rotationMatrix();
  Eigen::Vector3d t = pose_world_camera_.translation(); 
  // [R|t]
  Matrix34d Rt;
  Rt.leftCols<3>() = R;
  Rt.rightCols<1>() = t;

  // P = K * [R|t]
  Matrix34d P = K_ * Rt;
  return P;
    
  }
  

cv::Mat PerspectiveCameraModel::undistortImage(cv::Mat distorted_image) const
{
  // Convert to cv::Mats
  cv::Mat K_cv;
  cv::eigen2cv(K_, K_cv);
  cv::Mat dist_coeffs_cv;
  cv::eigen2cv(distortion_coeffs_, dist_coeffs_cv);

  // Undistort image.
  cv::Mat undistorted_image;
  cv::undistort(distorted_image, undistorted_image, K_cv, dist_coeffs_cv);

  return undistorted_image;
}
